import math
import random

# ================================
# 1. 计算两点欧氏距离
# ================================
def euclidean_distance(x1, x2):
    distance = 0
    for i in range(len(x1)):
        distance += (x1[i] - x2[i]) ** 2
    return math.sqrt(distance)

# ================================
# 2. KNN 分类器
# ================================
class KNN:
    def __init__(self, k=3):
        self.k = k
        self.X_train = None
        self.y_train = None

    # “训练”其实就是保存训练数据
    def fit(self, X, y):
        self.X_train = X
        self.y_train = y

    # 预测单个样本
    def predict_one(self, x):
        distances = []

        # 1）计算所有训练样本到 x 的距离
        for i in range(len(self.X_train)):
            d = euclidean_distance(x, self.X_train[i])
            distances.append((d, self.y_train[i]))   # (距离, 标签)

        # 2）按距离升序排序
        distances.sort(key=lambda t: t[0])

        # 3）取前 k 个
        k_neighbors = distances[:self.k]

        # 4）投票：统计出现次数最多的类别
        votes = {}
        for d, label in k_neighbors:
            votes[label] = votes.get(label, 0) + 1

        # 返回投票最多的类别
        return max(votes, key=votes.get)

    # 批量预测
    def predict(self, X):
        return [self.predict_one(x) for x in X]

# ================================
# 3. 测试代码（手动生成数据）
# ================================
if __name__ == "__main__":
    # 随机生成两类二维数据
    random.seed(0)

    X = []
    y = []

    # 类别0：中心在 (1,1)
    for _ in range(20):
        X.append([1 + random.random(), 1 + random.random()])
        y.append(0)

    # 类别1：中心在 (3,3)
    for _ in range(20):
        X.append([3 + random.random(), 3 + random.random()])
        y.append(1)

    # 创建模型
    knn = KNN(k=3)
    knn.fit(X, y)

    # 测试样本
    test = [[2, 2], [0.5, 0.5], [3.5, 3.2]]

    print("预测结果：", knn.predict(test))